# -*- coding: utf-8 -*-

# -------------------------------
# Author    : linlingwei
# Email     : 951035650@qq.com
# File Name : bodyCore.py
# Date      : 2018/11/6
# IDE       : PyCharm
# Version   : 1.1.1
# -------------------------------

import json

import maya.cmds as mc

import configs.constants.attribute_constant as ac
import configs.constants.nodeType_constant as nc
import configs.constants.suffix_constant as sc
import configs.constants.body_constant as bc
import configs.constants.global_constant as gc
reload(sc)
import scripts.commonCmds as con
import scripts.curveCmds as crc
import rig.jointTool as jt

__version__ = "1.2.3"

class BodyCore(object):
    """
    body rig base class.
    v1.1
    
    """
    def __init__(self,
                 localTag = "",
                 localMake = "",
                 localPrefix = ""):
        
        self.globalCtr = '_'.join(['global', sc.sCtr])
        self.globalStatic = con.lname(bc.gRig, sc.sStatic)
        self.transfer = "Transfer"
        
        self.localTag = localTag
        self.localMake = localMake
        self.localPrefix = localPrefix

        self.localName = con.lname(self.localPrefix,self.localMake+self.localTag)
        
        self.localGrp = con.lname(self.localPrefix,self.localTag,sc.sGrp)

        self.localPrefixGrp = con.lname(self.localName, sc.sGrp)
        
        self.localBridge = con.lname(self.localName,sc.sBridge)
        self.localStatic = None
        self.localSkinJotSet = None
        
        self.localSeq = list()
        
        self.follow = dict()
        
        self.local_colorId = bc.cLeft

        pass
    def refresh_localName(self):
        self.localName = con.lname(self.localPrefix, self.localMake + self.localTag)
    
        self.localGrp = con.lname(self.localPrefix, self.localTag, sc.sGrp)
    
        self.localPrefixGrp = con.lname(self.localName, sc.sGrp)
    
        self.localBridge = con.lname(self.localName, sc.sBridge)
        return

    def init(self):
        if not mc.objExists(bc.gBodyRig):
            bodyRig = mc.createNode(nc.tTransform,name=bc.gBodyRig)
            con.lockHideAttr(bodyRig,ignore=["v"])
            
            con.addStringAttr(bodyRig,"version",__version__)
            con.addStringAttr(bodyRig,"author",gc.gAuthor)
            
            mc.setAttr(bodyRig+".version", l=True)
            mc.setAttr(bodyRig+".author", l=True)
            mc.lockNode(bodyRig, lockName=True, lock=False)
    
        rigGrpName = "_".join([bc.gRig, sc.sGrp])
        if not mc.objExists(rigGrpName):
            rigGrp = mc.createNode(nc.tTransform,name=rigGrpName,p=bc.gBodyRig)
    
        rigStaticName = self.globalStatic
        if not mc.objExists(rigStaticName):
            rigStatic = mc.createNode(nc.tTransform,name=rigStaticName,p=bc.gBodyRig)
            con.lockHideAttr(rigStatic)
            mc.addAttr(rigStaticName, ln=ac.global_scale, at="float", dv=1, min=0.01, k=True)
            # mc.addAttr(rigStaticName,ln=gGlobalScale+"_offest",at="float",dv=1,min=0.01,k=True)

            mc.lockNode(rigStatic, lockName=True, lock=False)
        allCtrName = '_'.join(['all', sc.sCtr])
        if not mc.objExists(allCtrName):
            allCtr = crc.createControl(allCtrName, 'circle', 17)
            crc.scale(allCtr,10,10,10)
            allCtrGrp = con.emptyGroup(allCtr)
            con.lockHideAttr(allCtr, ['sx', 'sy', 'sz', 'v'])
            mc.parent(allCtrGrp[0],rigGrpName)
            
            gloCtr = crc.createControl(self.globalCtr, 'circle', 3)
            crc.scale(gloCtr,8,8,8)
            gCtrGrp = con.emptyGroup(gloCtr)
            mc.parent(gCtrGrp[0],allCtr)

            mc.addAttr(gloCtr, ln=ac.global_scale, at="float", dv=1, min=0.01, k=True)
            mc.addAttr(gloCtr, ln=ac.global_scale + "_offest", at="float", dv=1, min=0.01, k=False)
            mc.setAttr("{}.{}".format(gloCtr, ac.global_scale + "_offest"), l=True)

            gloMd = mc.createNode(nc.tMultiplyDivide, name = con.lname(ac.global_scale, sc.sMd))
            mc.connectAttr("{}.{}".format(gloCtr, ac.global_scale), "{}.{}".format(gloMd, "i1x"))
            mc.connectAttr("{}.{}".format(gloCtr, ac.global_scale + "_offest"), "{}.{}".format(gloMd, "i2x"))

            mc.connectAttr("{}.{}".format(gloMd, "ox"), "{}.{}".format(rigStaticName, ac.global_scale))

            for attr in ['sx', 'sy', 'sz']:
                mc.connectAttr("{}.{}".format(gloMd,"ox"),"{}.{}".format(gloCtr,attr))
            con.lockHideAttr(gloCtr, ['sx', 'sy', 'sz', 'v'])
        if not mc.objExists(self.transfer):
            mc.createNode(nc.tTransform,name=self.transfer,parent=self.globalCtr)
        if mc.objExists( bc.gBodyTemplate ):
            
            rel = mc.listRelatives(bc.gBodyTemplate,p=True)
            
            if not rel :
                if rel != bc.gBodyRig:
                    mc.lockNode(bc.gBodyTemplate,l=False)
                    mc.parent(bc.gBodyTemplate,bc.gBodyRig)
            mc.hide(bc.gBodyTemplate)
        
        if not mc.objExists(bc.gSkinJointSet):
            mc.sets(name = bc.gSkinJointSet,em=True)
          
        return rigGrpName, rigStaticName
    
    def localInit(self,static=False):
        
        self.init()
        
        localGrp = self.localGrp
        if not mc.objExists(localGrp):
            mc.createNode(nc.tTransform,name=localGrp,parent=self.globalCtr)
        
        localPrefixGrp = self.localPrefixGrp
        
        if not mc.objExists(localPrefixGrp):
            mc.createNode(nc.tTransform,name = localPrefixGrp,parent=localGrp)
        
        localBridge = self.localBridge
        if not mc.objExists(localBridge):
            mc.createNode(nc.tTransform,name = localBridge,parent=localPrefixGrp)
            if ac.global_scale not in mc.listAttr(self.localBridge):
                mc.addAttr(self.localBridge, ln=ac.global_scale, at='float', dv=1, k=True)
                mc.connectAttr(self.globalStatic +"." + ac.global_scale, self.localBridge + "." + ac.global_scale)
        if static:
            localStatic = con.lname(self.localPrefix,self.localName,sc.sStatic)
            self.localStatic = localStatic
            if not mc.objExists(localStatic):
                mc.createNode(nc.tTransform, name=localStatic, parent=self.globalStatic)

        localSkinJotSet = con.lname(self.localName,"skin_joint", sc.sSet)
        self.localSkinJotSet = localSkinJotSet
        if not mc.objExists(localSkinJotSet):
            mc.sets(name = localSkinJotSet,em=True)
            mc.sets(localSkinJotSet, include=bc.gSkinJointSet)
        return
    
    def parsePartTemplate(self):
        localSet = con.lname(self.localMake,self.localTag,bc.gTemplate,bc.sSet)
        if not mc.objExists(localSet):
            print localSet
            return False
        chSet = mc.sets(localSet, q=True)
        data = {}
        for cs in chSet:
            tag = mc.getAttr(cs + ".otherType")
            data[tag] = cs
        return data
    def getTemplateDate(self):
        localPrefix = con.lname(self.localMake,self.localTag, bc.gTemplate, bc.sSet)
        data = mc.getAttr(localPrefix + ".crvData")
        return json.loads(data)

    def followHook(self, tgtObj, fowAttr=(ac.global_scale, ac.local),onlyOrient=False,dv=1):
        
        flwHook = dict()
        
        objLoc = con.objToLocator(tgtObj,con.lname(tgtObj,sc.sFlw),parent=False)

        objGrp = con.emptyGroup(objLoc,["Zero","Con"])
        objCon = objGrp[1]
        flwlist = list()
        for flw in fowAttr:
            flwLoc = con.objToLocator(tgtObj, con.lname(tgtObj,flw), parent=False)
            flwloc_grp = con.emptyGroup(flwLoc,["Zero"])[0]
            flwlist.append(flwLoc)
            flwHook[flw] = flwloc_grp
        
        pot = mc.pointConstraint(flwlist,objCon,mo=True)[0]
        ort = mc.orientConstraint(flwlist,objCon,mo=True)[0]

        mc.addAttr(objCon, ln=bc.aFollow, at='enum', enumName=':'.join(fowAttr), dv=dv, k=True)

        for i,loc in enumerate(flwlist):
            cho = mc.createNode('choice', name=con.lname(tgtObj, fowAttr[i], sc.sCho))
            mc.connectAttr(objCon + ".{}".format(bc.aFollow), cho + ".selector")
            mc.connectAttr(cho + '.output', pot + '.%sW%s' % (loc,i))
            mc.connectAttr(cho + '.output', ort + '.%sW%s' % (loc,i))

            for idx in range(len(fowAttr)):
                if idx == i:
                    mc.setAttr(cho+'.input[%s]'%idx,1)
                else:
                    mc.setAttr(cho + '.input[%s]' % idx, 0)
            if onlyOrient:
                v = mc.getAttr(pot + '.%sW%s' % (loc, i))
                mc.disconnectAttr(cho + '.output', pot + '.%sW%s' % (loc, i))
                mc.setAttr(pot + '.%sW%s' % (loc, i),v)
        flwHook["out"] = objLoc
        flwHook[sc.sBridge] = objCon
        flwHook[bc.gRoot] = objGrp[0]

        return flwHook

    def ikfkSwitch(self, ikJotList, fkJotList, drJotlist,):

        if bc.aIkSwitch not in mc.listAttr(self.localBridge):
            mc.addAttr(self.localBridge, ln=bc.aIkSwitch, at='long', dv=0, max=1, min=0, k=False)
        if bc.aFkSwitch not in mc.listAttr(self.localBridge):
            mc.addAttr(self.localBridge, ln=bc.aFkSwitch, at='long', dv=0, max=1, min=0, k=False)
        if bc.aIkFkSwitch not in mc.listAttr(self.localBridge):
            mc.addAttr(self.localBridge, ln=bc.aIkFkSwitch, at='long', dv=0, max=1, min=0, k=True)
            ikfkRee = mc.createNode(nc.tReverse, name=con.lname(self.localName, bc.aIkFkSwitch, sc.sRee))
            mc.connectAttr(self.localBridge + '.%s' % bc.aIkFkSwitch, ikfkRee + '.inputX')
        
            mc.connectAttr(ikfkRee + '.outputX', self.localBridge + '.%s' % bc.aIkSwitch)
            mc.connectAttr(ikfkRee + '.inputX', self.localBridge + '.%s' % bc.aFkSwitch)
        
        for ij,fj,dj in zip(ikJotList, fkJotList, drJotlist):
        # for i in range(len(drJotlist)):
            ikfkPac = mc.parentConstraint(ij,fj,dj, mo=True)[0]
            mc.connectAttr(self.localBridge + '.%s' % bc.aIkSwitch, ikfkPac + '.%sW0' % ij)
            mc.connectAttr(self.localBridge + '.%s' % bc.aFkSwitch, ikfkPac + '.%sW1' % fj)
            
            # st = con.objToType(drJotlist[i],con.lname(drJotlist[i],bc.aLocalScale))
            # stp = mc.listRelatives(drJotlist[i],p=1)
            # if stp:
            #     mc.parent(st,stp[0])
            # ikfkSce = mc.scaleConstraint(ikJotList[i],
            #                               fkJotList[i],
            #                              st, mo=True)[0]
            # mc.connectAttr(st+".s",drJotlist[i]+".s")
            # mc.connectAttr(self.localBridge + '.%s' % bc.aIkSwitch, ikfkSce + '.%sW0' % ikJotList[i])
            # mc.connectAttr(self.localBridge + '.%s' % bc.aFkSwitch, ikfkSce + '.%sW1' % fkJotList[i])
            
        return
    
    def tangentTwist(self,objlist,name,twistSpan=3,endrot = None,nlist = []):
        ptlist = list()
        for j in objlist:
            pt = mc.xform(j, q=True, ws=True, t=True)
            ptlist.append(pt)
    
        crv = mc.curve(name=con.lname(name, sc.sCrv), d=1, p=ptlist)
        crvShape = con.findRelatedShape(crv)
        pro = 1./(len(objlist)-1)
        
        toPatNum = (twistSpan*(len(objlist)-1))-(len(objlist)-2)
        tploclist, tpgrplist= con.crvLoc(crv,toPatNum,name=con.lname(name,"toParent"),connect=False)
        
        tanobj = None
        sjotlist = list()

        for i,obj in enumerate(objlist[:-1]):
            
            
            mc.addAttr(obj,ln=bc.aLocalScale,at='float',dv=1,min=0.01,k=True)
            mc.addAttr(obj, ln=bc.aMidLocalScale, at='float', dv=1, min=0.01, k=True)
            
            tanIkObj1 = it.IkTool()
            tanjotlist1, tanloclist1 = tanIkObj1.tangentSolve(objlist[i], objlist[i+1],
                                                              con.lname(name, str(i)),)

            mc.parent(tanloclist1[0],tploclist[i*2]) #0 2 4
            mc.parent(tanloclist1[1],tploclist[i*2+1])#1 3 5
            mc.parent(tanloclist1[2],tploclist[i*2+2])#2 4 6
            
            mc.connectAttr(obj+"."+bc.aLocalScale,tanIkObj1.localBridge+"."+bc.aStarLocalScale)
            mc.connectAttr(obj + "." + bc.aMidLocalScale, tanIkObj1.localBridge + "." + bc.aMidLocalScale)
            if tanobj:
                mc.connectAttr(tanjotlist1[0]+"."+bc.aRotation,tanobj.localBridge+"."+bc.aEndRotation)
                mc.connectAttr(obj + "." + bc.aLocalScale, tanobj.localBridge + "." + bc.aEndLocalScale)

            tanobj = tanIkObj1
            mc.parent(tanIkObj1.localGrp,self.localGrp)
            mc.parent(tanIkObj1.localStatic,self.localStatic)
            
            mc.sets(tanjotlist1,include=self.localSkinJotSet)
            
            if nlist:
                lpat = None
                llist = []
                for x,tj in enumerate(tanjotlist1[:-1]):
                    sj = mc.createNode(nc.tJoint,name=con.lname(self.localPrefix,nlist[i]+str(x+1)))
                    mc.delete(mc.pointConstraint(tj,sj,mo=False))
                    if lpat:
                        mc.parent(sj,lpat)
                    # mc.parent(mc.parentConstraint(tj,sj,mo=True),self.localStatic)
                    # mc.parent(mc.scaleConstraint(tj,sj,mo=True),self.localStatic)
                    lpat = sj
                    llist.append(sj)
                    
                if sjotlist:
                    mc.parent(llist[0],sjotlist[-1])
                sjotlist+=llist
            
                mc.joint(llist, e=True, oj='yzx', secondaryAxisOrient='zup', ch=True, zso=True)
                mc.setAttr('%s.jointOrient' % (llist[-1]), 0, 0, 0)
                # spat = None
                for tj,sj in zip(tanjotlist1[:-1],llist):
                    mc.parent(mc.parentConstraint(tj,sj,mo=True),self.localStatic)
                    # mc.connectAttr(tj+".s",sj+".s")
                st = self.scaleTransfer(tj,sj)
                # if spat:
                #     mc.parent(st,spat)
                # spat = st
                #
        mc.addAttr(objlist[-1], ln=bc.aLocalScale, at='float', dv=1, min=0.01, k=True)
        mc.connectAttr(objlist[-1] + "." + bc.aLocalScale, tanobj.localBridge + "." + bc.aEndLocalScale)

        for i,x in enumerate(range(0,toPatNum,2)):
            outlocShape = con.findRelatedShape(tploclist[x])
            mc.connectAttr(outlocShape+".wp[0]",crvShape+".cp[%s]"%i)
        
        if endrot:
            mc.connectAttr(endrot,tanobj.localBridge+"."+bc.aEndRotation)
        
        mc.hide(tpgrplist)
        mc.parent(crv,self.localStatic)
        
        return tploclist, tpgrplist,sjotlist
    
    def localJoint(self,tmpData, tag='',localSeq = [], mirror = False):
        
        
        tmpData[bc.gSeq] = localSeq
        tjotlist = jt.seqJointChain(tmpData, con.lname("ltmp_",self.localName,tag,bc.gTemplate))
        if mirror:
            mjotlist = mc.mirrorJoint(tjotlist[0], mirrorYZ=True, mirrorBehavior=True,
                                      searchReplace=("ltmp__"  , "rtmp__" ))
            mc.delete(tjotlist)
            tjotlist = mjotlist
        # mc.hide(tjotlist[0])
        for i, j in enumerate(tjotlist):
            nn = j.replace("_" + bc.gTemplate, "").replace("ltmp__","").replace("rtmp__","")
            tjotlist[i] = mc.rename(j, nn)

        return tjotlist
    
    def addTwist(self,twistInJotlist,twistEndJotlist,twistTmp,mirror=False):
    
        # twistInJotlist = drJotlist[1:4]
        nrjot = jt.noroll(twistEndJotlist[0], twistEndJotlist[1], con.lname(self.localName, bc.gWrist))[0]
    
        tploclist, tpgrplist,sjotlist = self.tangentTwist(twistInJotlist, con.lname(self.localName, bc.gSec),
                                                 endrot=nrjot + ".ry")

        secCtrllist = list()
        
        mc.addAttr(self.localBridge, ln=bc.aSecVis, at="long", dv=1, min=0, max=1, k=True)
        # twistTmp = [tmp[gArm], tmp[gForeArm], tmp[gWrist]]
        for i, tl in enumerate(tpgrplist):
            secCtrl = con.objToControl(twistTmp[i / 2], con.lname( self.localName, "Twist" + str(i)),
                                         mirror=mirror, colorId=bc.cSec, scale=0.8)
            mc.delete(mc.pointConstraint(tl, secCtrl, mo=False))
            mc.parent(tl, secCtrl)
            grp = con.emptyGroup(secCtrl)
            mc.parent(grp[0], twistInJotlist[i / 2])

            mc.connectAttr(self.localBridge + "." + bc.aSecVis, grp[0] + ".v")

            secCtrllist.append(secCtrl)

        mc.pointConstraint(secCtrllist[0], secCtrllist[2], secCtrllist[1] + "_Con")
        mc.pointConstraint(secCtrllist[2], secCtrllist[4], secCtrllist[3] + "_Con")
        
        return sjotlist
    def scaleTransfer(self,sobj,tobj):
        ttr = mc.createNode(nc.tTransform, name=con.lname(tobj, self.transfer),
                            parent=self.transfer)
        mc.scaleConstraint(sobj, ttr, mo=True)
        mc.connectAttr(ttr + ".s", tobj + ".s")
        return ttr
    def transfers(self,sobj,tobj):
        
        ttr = mc.createNode(nc.tTransform, name=con.lname(tobj, self.transfer),
                            parent=self.transfer)
        # pt = mc.listRelatives(tobj,p=1)
        # mc.parent(ttr,cCore.lname(pt, self.transfer))
        mc.delete(mc.parentConstraint(tobj, ttr, mo=False))
        mc.parentConstraint(sobj, ttr, mo=True)
        mc.scaleConstraint(sobj, ttr, mo=True)

        mc.parentConstraint(ttr, tobj, mo=True)

        mc.connectAttr(ttr + ".s", tobj + ".s")
        return ttr

    def addLocalScale(self,*args):
        for arg in args:
            mc.addAttr(arg, ln=ac.local_scale, at='float', dv=1, min=0.01, k=True)
            mc.connectAttr(arg + "." + ac.local_scale, arg + ".sx")
            mc.connectAttr(arg + "." + ac.local_scale, arg + ".sy")
            mc.connectAttr(arg + "." + ac.local_scale, arg + ".sz")
        
            con.lockHideAttr(arg, ["sx", 'sy', 'sz'])
    
        return

    def seqJointChain(self,seqDict,seq, name):
    
        # seqs = seqDict[bc.gSeq]
        jotlist = list()
        for s in seq:
            jot = jt.joint(seqDict[s], name=con.lname(name,s,sc.sJot))
            jotlist.append(jot)
        jt.jointChain(jotlist,axial="x")
    
        return jotlist
    def create_joint(self,tag,seq,mirror=False):
        tmpData = self.parsePartTemplate()
        # for tmp in tmpData:
        jotlist = self.seqJointChain(tmpData,seq, con.lname(self.localName,tag))
        if mirror:
            mdrJotlist = mc.mirrorJoint(jotlist[0], mirrorYZ=True, mirrorBehavior=True)
            mc.delete(jotlist)
            [mc.rename(m,j) for m,j in zip(mdrJotlist,jotlist)]
            # jotlist=mdrJotlist
        return jotlist
if __name__ == '__main__':
    pass